#!/bin/bash

# DeepGuard 安全检测系统启动脚本
# 支持 macOS、Ubuntu、CentOS 等多种操作系统

set -e  # 遇到错误立即退出

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# 日志函数
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检测操作系统和架构
detect_os() {
    local arch=$(uname -m)
    
    if [[ "$OSTYPE" == "darwin"* ]]; then
        OS="macos"
        DISTRO="macOS"
        PLATFORM="mac"
        COMPOSE_FILE="docker-compose.mac.yml"
    elif [[ -f /etc/os-release ]]; then
        . /etc/os-release
        OS="linux"
        DISTRO=$ID
        PLATFORM="linux"
        COMPOSE_FILE="docker-compose.linux.yml"
    elif [[ -f /etc/redhat-release ]]; then
        OS="linux"
        DISTRO="centos"
        PLATFORM="linux"
        COMPOSE_FILE="docker-compose.linux.yml"
    else
        OS="unknown"
        DISTRO="unknown"
        PLATFORM="unknown"
        COMPOSE_FILE="docker-compose.yml"
    fi
    
    log_info "检测到操作系统: $DISTRO ($arch)"
    log_info "使用配置文件: $COMPOSE_FILE"
    
    # 检查配置文件是否存在
    if [[ ! -f "deep-guard-docker/$COMPOSE_FILE" ]]; then
        log_warning "配置文件不存在: $COMPOSE_FILE，使用默认配置"
        COMPOSE_FILE="docker-compose.yml"
    fi
}

# 检查命令是否存在
command_exists() {
    command -v "$1" >/dev/null 2>&1
}

# 安装 Docker (Linux)
install_docker_linux() {
    log_info "正在安装 Docker..."
    
    case $DISTRO in
        "ubuntu"|"debian")
            sudo apt-get update
            sudo apt-get install -y apt-transport-https ca-certificates curl gnupg lsb-release
            curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /usr/share/keyrings/docker-archive-keyring.gpg
            echo "deb [arch=$(dpkg --print-architecture) signed-by=/usr/share/keyrings/docker-archive-keyring.gpg] https://download.docker.com/linux/ubuntu $(lsb_release -cs) stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null
            sudo apt-get update
            sudo apt-get install -y docker-ce docker-ce-cli containerd.io docker-compose-plugin
            ;;
        "centos"|"rhel"|"fedora")
            sudo yum install -y yum-utils
            sudo yum-config-manager --add-repo https://download.docker.com/linux/centos/docker-ce.repo
            sudo yum install -y docker-ce docker-ce-cli containerd.io docker-compose-plugin
            sudo systemctl start docker
            sudo systemctl enable docker
            ;;
        "arch")
            sudo pacman -S docker docker-compose --noconfirm
            sudo systemctl start docker
            sudo systemctl enable docker
            ;;
        *)
            log_error "不支持的 Linux 发行版: $DISTRO"
            log_info "请手动安装 Docker 和 Docker Compose"
            exit 1
            ;;
    esac
    
    # 添加当前用户到 docker 组
    sudo usermod -aG docker $USER
    log_warning "请注销并重新登录以使 Docker 组权限生效，或运行: newgrp docker"
}

# 安装 Docker (macOS)
install_docker_macos() {
    log_info "正在安装 Docker Desktop for Mac..."
    
    if command_exists brew; then
        brew install --cask docker
    else
        log_error "未检测到 Homebrew，请手动安装 Docker Desktop"
        log_info "下载地址: https://www.docker.com/products/docker-desktop"
        exit 1
    fi
}

# 检查并安装 Docker
check_docker() {
    log_info "检查 Docker 安装状态..."
    
    if ! command_exists docker; then
        log_warning "Docker 未安装，正在自动安装..."
        case $OS in
            "macos")
                install_docker_macos
                ;;
            "linux")
                install_docker_linux
                ;;
            *)
                log_error "不支持的操作系统，请手动安装 Docker"
                exit 1
                ;;
        esac
    else
        log_success "Docker 已安装"
    fi
    
    # 检查 Docker 是否运行
    if ! docker info >/dev/null 2>&1; then
        log_warning "Docker 服务未运行，正在启动..."
        case $OS in
            "macos")
                open -a Docker
                log_info "请等待 Docker Desktop 启动完成..."
                sleep 10
                ;;
            "linux")
                sudo systemctl start docker
                ;;
        esac
        
        # 等待 Docker 启动
        local count=0
        while ! docker info >/dev/null 2>&1; do
            if [ $count -ge 30 ]; then
                log_error "Docker 启动超时，请检查 Docker 安装"
                exit 1
            fi
            echo -n "."
            sleep 2
            ((count++))
        done
        echo ""
        log_success "Docker 服务已启动"
    fi
}

# 检查并安装 Docker Compose
check_docker_compose() {
    log_info "检查 Docker Compose 安装状态..."
    
    # 优先检查 Docker Compose V2 (docker compose) - 新版本
    if docker compose version >/dev/null 2>&1; then
        COMPOSE_CMD="docker compose"
        log_success "Docker Compose V2 已安装 (推荐版本)"
        return
    fi
    
    # 检查 Docker Compose V1 (docker-compose) - 兼容旧版本
    if command_exists docker-compose; then
        COMPOSE_CMD="docker-compose"
        log_success "Docker Compose V1 已安装 (兼容模式)"
        log_warning "建议升级到 Docker Compose V2 (docker compose)"
        return
    fi
    
    # Docker Compose V2 通常随 Docker Desktop 一起安装
    log_error "Docker Compose 未检测到"
    log_info "Docker Compose V2 通常随 Docker Desktop 一起安装"
    log_info "请确保 Docker Desktop 已正确安装并启动"
    
    case $OS in
        "macos")
            log_info "Mac用户: 请安装 Docker Desktop for Mac"
            log_info "下载地址: https://www.docker.com/products/docker-desktop"
            ;;
        "linux")
            log_info "Linux用户: 请安装 Docker Engine 和 Docker Compose Plugin"
            log_info "Ubuntu/Debian: sudo apt-get install docker-compose-plugin"
            log_info "CentOS/RHEL: sudo yum install docker-compose-plugin"
            ;;
    esac
    
    exit 1
}

# 检查环境变量配置
check_env_config() {
    log_info "检查环境变量配置..."
    
    if [[ ! -f ".env" ]]; then
        log_warning ".env 文件不存在，正在从 .env_example 创建..."
        cp .env_example .env
        log_info "请编辑 .env 文件，配置您的SMTP邮件服务信息"
        return 1
    fi
    
    # 检查必要的环境变量
    local missing_vars=()
    
    if ! grep -q "^SMTP_HOST=" .env || grep -q "^SMTP_HOST=$" .env; then
        missing_vars+=("SMTP_HOST")
    fi
    
    if ! grep -q "^SMTP_PORT=" .env || grep -q "^SMTP_PORT=$" .env; then
        missing_vars+=("SMTP_PORT")
    fi
    
    if ! grep -q "^SMTP_USER=" .env || grep -q "^SMTP_USER=$" .env; then
        missing_vars+=("SMTP_USER")
    fi
    
    if ! grep -q "^SMTP_PASSWORD=" .env || grep -q "^SMTP_PASSWORD=$" .env; then
        missing_vars+=("SMTP_PASSWORD")
    fi

    if ! grep -q "^SMTP_SENDER_EMAIL=" .env || grep -q "^SMTP_SENDER_EMAIL=$" .env; then
        missing_vars+=("SMTP_SENDER_EMAIL")
    fi

    if ! grep -q "^DEFAULT_RECIPIENT_EMAIL=" .env || grep -q "^DEFAULT_RECIPIENT_EMAIL=$" .env; then
        missing_vars+=("DEFAULT_RECIPIENT_EMAIL")
    fi
    
    if [ ${#missing_vars[@]} -gt 0 ]; then
        log_warning "以下SMTP环境变量未配置: ${missing_vars[*]}"
        log_info "告警功能可能无法正常工作，请编辑 .env 文件进行配置"
        return 1
    fi
    
    log_success "环境变量配置完整"
    return 0
}

# 检查端口占用
check_ports() {
    log_info "检查端口占用情况..."
    
    local ports=(23847 19563 16842)
    local occupied_ports=()
    
    for port in "${ports[@]}"; do
        if command_exists lsof; then
            if lsof -i :$port >/dev/null 2>&1; then
                occupied_ports+=($port)
            fi
        elif command_exists netstat; then
            if netstat -tuln | grep ":$port " >/dev/null 2>&1; then
                occupied_ports+=($port)
            fi
        elif command_exists ss; then
            if ss -tuln | grep ":$port " >/dev/null 2>&1; then
                occupied_ports+=($port)
            fi
        fi
    done
    
    if [ ${#occupied_ports[@]} -gt 0 ]; then
        log_warning "以下端口被占用: ${occupied_ports[*]}"
        log_info "如果是 DeepGuard 服务占用，将会自动重启"
    else
        log_success "所有端口可用"
    fi
}

# 构建和启动服务
start_services() {
    log_info "正在构建和启动 DeepGuard 服务..."
    log_info "平台: $PLATFORM, 配置文件: $COMPOSE_FILE"
    
    cd deep-guard-docker
    
    # 停止现有服务（尝试所有可能的配置文件）
    log_info "停止现有服务..."
    for config in docker-compose.yml docker-compose.mac.yml docker-compose.linux.yml; do
        if [[ -f "$config" ]]; then
            $COMPOSE_CMD -f "$config" down --remove-orphans 2>/dev/null || true
        fi
    done
    
    # 启用Docker BuildKit加速构建
    export DOCKER_BUILDKIT=1
    export COMPOSE_DOCKER_CLI_BUILD=1
    log_info "已启用 Docker BuildKit 加速构建"
    
    # 预下载基础镜像
    log_info "预下载基础镜像..."
    docker pull python:3.9-slim >/dev/null 2>&1 &
    docker pull redis:7-alpine >/dev/null 2>&1 &
    
    # 根据平台预下载ClamAV镜像
    if [[ "$PLATFORM" == "mac" ]]; then
        log_info "预下载Mac兼容的ClamAV镜像..."
        docker pull clamav/clamav:1.0.0_base >/dev/null 2>&1 &
    else
        log_info "预下载Linux优化的ClamAV镜像..."
        docker pull clamav/clamav:stable_base >/dev/null 2>&1 &
    fi
    wait
    
    # 构建镜像（使用清华源加速）
    log_info "构建 Docker 镜像（使用清华源加速）..."
    $COMPOSE_CMD -f "$COMPOSE_FILE" build --no-cache --parallel
    
    # 启动服务
    log_info "启动服务容器..."
    $COMPOSE_CMD -f "$COMPOSE_FILE" up -d
    
    cd ..
    
    # 等待服务启动
    log_info "等待服务启动..."
    sleep 15
    
    # 检查服务状态
    local services=("deepguard-alarm:23847" "deepguard-scanner:19563" "deepguard-redis:16842")
    local failed_services=()
    
    for service in "${services[@]}"; do
        local name=${service%:*}
        local port=${service#*:}
        
        if ! docker ps | grep -q $name; then
            failed_services+=($name)
        fi
    done
    
    if [ ${#failed_services[@]} -gt 0 ]; then
        log_error "以下服务启动失败: ${failed_services[*]}"
        log_info "查看日志: cd deep-guard-docker && $COMPOSE_CMD -f $COMPOSE_FILE logs"
        return 1
    fi
    
    log_success "所有服务启动成功"
}

# 显示服务状态
show_status() {
    log_info "服务状态:"
    echo ""
    
    cd deep-guard-docker
    $COMPOSE_CMD -f "$COMPOSE_FILE" ps
    cd ..
    
    echo ""
    log_info "服务访问地址:"
    echo "  📊 告警服务:   http://localhost:23847"
    echo "  🔍 扫描服务:   http://localhost:19563"
    echo "  💾 Redis数据库: localhost:16842"
    echo "  🦠 ClamAV引擎: localhost:3310"
    echo ""
    
    if [[ "$PLATFORM" == "mac" ]]; then
        log_info "Mac平台 - 使用兼容ARM64的ClamAV镜像"
    else
        log_info "Linux平台 - 使用内存优化的ClamAV镜像 (限制1GB内存)"
    fi
    
    echo ""
    log_info "运行测试: ./test.sh"
    log_info "查看日志: cd deep-guard-docker && $COMPOSE_CMD -f $COMPOSE_FILE logs -f"
    log_info "停止服务: cd deep-guard-docker && $COMPOSE_CMD -f $COMPOSE_FILE down"
}

# 主函数
main() {
    echo "🚀 DeepGuard 安全检测系统启动脚本"
    echo "===================================="
    echo ""
    
    # 检测操作系统
    detect_os
    
    # 检查 Docker
    check_docker
    
    # 检查 Docker Compose
    check_docker_compose
    
    # 检查环境配置
    env_ok=true
    check_env_config || env_ok=false
    
    # 检查端口
    check_ports
    
    # 启动服务
    if start_services; then
        echo ""
        log_success "🎉 DeepGuard 系统启动成功!"
        echo ""
        show_status
        
        if [ "$env_ok" = false ]; then
            echo ""
            log_warning "⚠️  请配置 .env 文件中的SMTP邮件服务信息以启用告警功能"
        fi
    else
        log_error "❌ 系统启动失败"
        exit 1
    fi
}

# 执行主函数
main "$@"